import cv2
import numpy as np
from ultralytics import YOLO

CLASSES = ['number']

COLORS = np.random.uniform(0, 255, size=(len(CLASSES), 3))

def draw_boxes(image, results):
    mask_img = image.copy()
    num_boxes = []
    for box in results[0].boxes:
        cls = int(box.cls.cpu().numpy()[0])
        color = COLORS[cls]
        xyxy = box.xyxy[0].cpu().numpy().tolist()
        x1, y1, x2, y2 = [int(_) for _  in xyxy]
        cv2.rectangle(mask_img, (x1, y1), (x2, y2), color, 2)
        box = mask_img[y1:y2, x1:x2, :]
        num_boxes.append(box)
    return mask_img, num_boxes

model_path = 'models/yolov11_number_best_v1.onnx'
number_detector = YOLO(model_path, task='detect') 
print("===== init model ok =======")

image_path = 'datas/v000_1.jpg'
image = cv2.imread(image_path)

results = number_detector(image)
print("======= detect results ", results)

result_img, boxes = draw_boxes(image, results)
cv2.imwrite("detect_number.jpg", result_img)

for idx, box in enumerate(boxes):
    cv2.imwrite("number_{}.jpg".format(idx), box)

